package alfred.mdp.broker;

import alfred.mdp.Majordomo;
import alfred.mdp.messages.ProtocolMessage;
import alfred.mdp.messages.client.Reply;
import alfred.mdp.messages.client.Request;
import alfred.mdp.messages.worker.Ready;
import alfred.mdp.utils.Shutdowner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.ZContext;
import org.zeromq.ZDispatcher;
import org.zeromq.ZFrame;
import org.zeromq.ZMQ;
import org.zeromq.ZMsg;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;

import static alfred.mdp.messages.ProtocolMessage.get;
import static alfred.mdp.messages.ProtocolMessage.getSubMsg;
import static alfred.mdp.utils.Utils.asString;

/**
 */
public class Broker {

    private static final Logger LOG = LoggerFactory.getLogger(Broker.class);

    private final String id;
    private final ExecutorService scheduler = Executors.newCachedThreadPool();
    private final ConcurrentMap<String, List<WorkerReference>> workersForService = new ConcurrentHashMap<String, List<WorkerReference>>();
    private final ConcurrentMap<String, BlockingQueue<Request>> requestsForService = new ConcurrentHashMap<String, BlockingQueue<Request>>();
    private final ZMQ.Socket clientSocket;
    private final ZMQ.Socket workerSocket;
    private final ZDispatcher dispatcher;
    private final ZDispatcher.ZSender clientSender;
    private final ZDispatcher.ZSender workerSender;

    public Broker(ZContext context, ZDispatcher dispatcher, String clientEndpoint, String workerEndpoint) {
        this.id = UUID.randomUUID().toString();
        this.clientSocket = context.createSocket(ZMQ.ROUTER);
        this.workerSocket = context.createSocket(ZMQ.ROUTER);
        this.clientSocket.bind(clientEndpoint);
        this.workerSocket.bind(workerEndpoint);
        this.dispatcher = dispatcher;
        this.clientSender = new ZDispatcher.ZSender();
        this.workerSender = new ZDispatcher.ZSender();
        dispatcher.registerHandler(clientSocket, new ClientHandler(), clientSender);
        dispatcher.registerHandler(workerSocket, new WorkerHandler(), workerSender);
        Shutdowner.add(getShutdownGroup(), scheduler.submit(new CheckWorkersTask(this)));
    }

    public void registerWorker(ByteBuffer workerReferenceId, String serviceName) {
        LOG.trace("registerWorker");
        if (getWorkerReference(workerReferenceId) == null) {
            LOG.debug("registerWorker for workerId {} and service name {}", asString(workerReferenceId), serviceName);
            requestsForService.putIfAbsent(serviceName, new LinkedBlockingQueue<Request>());
            WorkerReference workerReference = new WorkerReference(this, serviceName, workerReferenceId);
            putOn(workersForService, serviceName, workerReference);
            Shutdowner.add(workerReference.getShutdownGroup(), scheduler.submit(workerReference));
        }
    }

    public void unregisterWorker(WorkerReference workerReference, boolean sendDisconnect) {
        LOG.trace("unregisterWorker");
        if (workerReference != null) {
            LOG.debug("unregisterWorker for workerId {} and service name {}", asString(workerReference.getWorkerId()), workerReference.getServiceName());
            workerReference.shutdown(sendDisconnect);
            workersForService.get(workerReference.getServiceName()).remove(workerReference);
        }
    }

    public void shutdown() {
        LOG.trace("shutdown");
        Shutdowner.shutdown(getShutdownGroup());
        for (WorkerReference workerReference : valuesOf(workersForService)) {
            unregisterWorker(workerReference, true);
        }
        dispatcher.unregisterHandler(clientSocket);
        dispatcher.unregisterHandler(workerSocket);
        LOG.info("shutdown completed");
    }

    public void onClientRequest(ZMsg msg) {
        LOG.trace("onClientRequest");
        Request received = new Request(msg);
        putOnQueue(requestsForService, received.getServiceName(), received);
    }

    public void onHeartbeat(ByteBuffer workerReferenceId) {
        if (LOG.isTraceEnabled()) {
            LOG.trace("onHeartbeat for workerId {}", asString(workerReferenceId));
        }
        WorkerReference workerReference = getWorkerReference(workerReferenceId);
        if (workerReference != null && workerReference.isAlive()) {
            workerReference.onHeartbeat();

        }
    }

    public void checkWorkers() {
        if (LOG.isTraceEnabled()) {
            LOG.trace("checkWorkers");
        }
        for (WorkerReference workerReference : valuesOf(workersForService)) {
            if (!workerReference.isAlive()) {
                unregisterWorker(workerReference, true);
            }
        }
    }

    public WorkerReference getWorkerReference(ByteBuffer workerReferenceId) {
        for (WorkerReference workerReference : valuesOf(workersForService)) {
            if (workerReference.getWorkerId().equals(workerReferenceId)) {
                return workerReference;
            }
        }
        return null;
    }

    public void sendClientResponse(ZMsg received, ByteBuffer workerReferenceId) {
        if (LOG.isTraceEnabled()) {
            LOG.trace("sendClientResponse for workerId {}", asString(workerReferenceId));
        }
        WorkerReference workerReference = getWorkerReference(workerReferenceId);

        if (workerReference != null) {
            ZFrame clientId = get(received, 4);
            ZMsg response = getSubMsg(received, 6);
            clientSender.send(new Reply(clientId, workerReference.getServiceName(), response).getZMsg());
            if (LOG.isDebugEnabled()) {
                LOG.debug("sent ClientResponse for workerId {} to clientId {}", asString(workerReferenceId), asString(clientId.getData()));
            }

        } else {
            unregisterWorker(workerReference, true);

        }
    }

    public void sendMessageToWorker(ProtocolMessage message, byte[] workerId) {
        ZMsg msg = new ZMsg();
        msg.addLast(new ZFrame(workerId));
        msg.addAll(message.getZMsg());
        workerSender.send(msg);
    }

    public String getShutdownGroup() {
        return getClass().getName() + id;
    }

    public ConcurrentMap<String, BlockingQueue<Request>> getRequestsForService() {
        return requestsForService;
    }

    public ExecutorService getExecutor() {
        return scheduler;
    }

    public String getId() {
        return id;
    }

    private static <K, V> boolean putOn(ConcurrentMap<K, List<V>> map, K key, V value) {
        map.putIfAbsent(key, new ArrayList<V>());
        return map.get(key).add(value);
    }

    private static <K, V> List<V> valuesOf(ConcurrentMap<K, List<V>> map) {
        ArrayList<V> result = new ArrayList<V>();
        for (List<V> values : map.values()) {
            result.addAll(values);
        }
        return result;
    }

    private static <K, V> boolean putOnQueue(ConcurrentMap<K, BlockingQueue<V>> map, K key, V value) {
        map.putIfAbsent(key, new LinkedBlockingQueue<V>());
        return map.get(key).add(value);
    }

    private class ClientHandler implements ZDispatcher.ZMessageHandler {

        @Override
        public void handleMessage(ZDispatcher.ZSender sender, ZMsg received) {
            onClientRequest(received);
        }
    }

    private class WorkerHandler implements ZDispatcher.ZMessageHandler {

        @Override
        public void handleMessage(ZDispatcher.ZSender sender, ZMsg received) {
            ByteBuffer workerReferenceId = ByteBuffer.wrap(received.getFirst().getData());

            ZFrame type = get(received, 3);
            if (type.equals(Majordomo.Worker.HEARTBEAT)) {
                onHeartbeat(workerReferenceId);

            } else if (type.equals(Majordomo.Worker.DISCONNECT)) {
                unregisterWorker(getWorkerReference(workerReferenceId), false);

            } else if (type.equals(Majordomo.Worker.READY)) {
                registerWorker(workerReferenceId, new Ready(received).getServiceName());

            } else if (type.equals(Majordomo.Worker.REPLY)) {
                sendClientResponse(received, workerReferenceId);

            } else {
                unregisterWorker(getWorkerReference(workerReferenceId), true);

            }
        }
    }
}
